module partition_of_unity

    use space_manager
    use dynamic_quantities, only: &
        nq => num_quantity,       &
        np => num_parcel,         &
        x => location,            &
        f => quantities,          &
        quantities_name
    use tree_struct
    use data_fit
    use message
    use run_manager

    implicit none

    integer :: approx_order
    real :: critical_error
    real :: exp_factor
    character(20) :: approximate_method
    character(20) :: weight_function_type

    namelist /partition_of_unity_parameters/ &
        approximate_method, approx_order,    &
        weight_function_type, exp_factor,    &
        critical_error

    interface partition_of_unity_evaluate
        module procedure partition_of_unity_evaluate_single
        module procedure partition_of_unity_evaluate_batch
    end interface partition_of_unity_evaluate

    interface partition_of_unity_derivative
        module procedure partition_of_unity_derivative_single
        module procedure partition_of_unity_derivative_batch
    end interface partition_of_unity_derivative

contains

    subroutine partition_of_unity_init(namelist_file)
        character(*) :: namelist_file

        ! default value
        critical_error = 1.0E-4
        weight_function_type = "exponential"

        open(10, file=namelist_file)
        read(10, nml=partition_of_unity_parameters, end=10)
10      close(10)

        ! initialize the needed module
        call tree_struct_init(namelist_file, nq, approx_order)
        call data_fit_init(approximate_method, num_dim, approx_order)

        return
    end subroutine partition_of_unity_init

    subroutine partition_of_unity_finalize

        call tree_struct_finalize
        call data_fit_finalize

        return
    end subroutine partition_of_unity_finalize

    subroutine partition_of_unity_approximate(qi)
        integer, intent(in) :: qi ! quantity index

        integer :: n, res
        integer :: i, j, k

        ! internal control variable, leave them alone
        integer, save :: counter = 1
        character(50) :: file_name

        if(counter == 1) then
            call tree_struct_reset
            call tree_struct_register_data(np, x)
        end if
        if(counter == nq) then
            counter = 1
        else
            counter = counter+1
        end if

        do k = 1, num_tree(3)
        do j = 1, num_tree(2)
        do i = 1, num_tree(1)
            n = forest(qi)%tree(i,j,k)%np
            res = subdomain_approximate(qi, forest(qi)%tree(i,j,k))
        end do
        end do
        end do

        return
    end subroutine partition_of_unity_approximate

    recursive integer function subdomain_approximate(qi, node) &
        result(res)
        integer, intent(in) :: qi
        type(tree_node), intent(inout), target :: node

        integer :: i, j, k, n, ret
        real, allocatable :: xx(:,:), ff(:)
        real, allocatable :: f0(:)
        real :: error0

        n = node%np

        allocate(xx(num_dim,n))
        allocate(ff(n))
        allocate(f0(n))

        do j = 1, n
            do i = 1, num_dim
                xx(i,j) = x(i,node%pi(j))+ &
                    (boundary(2,i)-boundary(1,i))*node%pb(i,j)
            end do
        end do

        ff = f(node%pi(1:n), qi)

        call data_fit_approximate(n, xx, ff, node%approx_order, node%c)
        call data_fit_evaluate(n, xx, f0, node%approx_order,  node%c)

        error0 = 0.0
        do i = 1, n
            error0 = error0+(ff(i)-f0(i))**2
        end do
        error0 = sqrt(error0/n)

        if(error0 <= critical_error) then
            res = 1
        else if(node%lev >= max_lev) then
            res = 1
        else if(node%np <= data_fit_min_num_data(node%approx_order)) then
            res = 1
        else
            res = 2
            call tree_struct_subdivide_tree(node, np, x)
            do k = 1, node%num_child(3)
                do j = 1, node%num_child(2)
                    do i = 1, node%num_child(1)
                        n = node%child(i,j,k)%np
                        ret = subdomain_approximate(qi, node%child(i,j,k))
                    end do
                end do
            end do
        end if

        return
    end function subdomain_approximate

    subroutine partition_of_unity_evaluate_batch(qi, n, x1, f1)
        integer, intent(in) :: qi, n
        real, intent(in) :: x1(num_dim,n)
        real, intent(out) :: f1(n)

        integer :: i

        do i = 1, n
            call partition_of_unity_evaluate_single(qi, x1(:,i), f1(i))
        end do

        return
    end subroutine partition_of_unity_evaluate_batch

    subroutine partition_of_unity_evaluate_single(qi, x1, f1)
        integer, intent(in) :: qi
        real, intent(in) :: x1(num_dim)
        real, intent(out) :: f1

        integer :: or(3) ! reference root tree index
        integer :: n(3)  ! evaluate tree number
        integer :: o(3)
        integer :: i, j, k, res
        real :: sum_weight

        ! root tree index where point resides
        call tree_struct_where_point(x1, or)

        ! set the number of surrounding root tree
        ! ---> needs careful consideration
        !      it should depend on the weight function
        n = 0
        n(1:num_dim) = 1

        f1 = 0.0
        sum_weight = 0.0
        do k = or(3)-n(3), or(3)+n(3)
        do j = or(2)-n(2), or(2)+n(2)
        do i = or(1)-n(1), or(1)+n(1)
            call tree_struct_boundary_check((/i,j,k/), o)
            res = subdomain_evaluate(qi,         &
                forest(qi)%tree(o(1),o(2),o(3)), &
                x1, f1, sum_weight)
        end do
        end do
        end do
        f1 = f1/sum_weight

        return
    end subroutine partition_of_unity_evaluate_single

    recursive integer function subdomain_evaluate(qi, node, x1, f1, sum_weight) result(res)
        integer, intent(in) :: qi
        type(tree_node), intent(in) :: node
        real, intent(in) :: x1(num_dim)
        real, intent(out) :: f1
        real, intent(out) :: sum_weight

        integer :: i, j, k
        real :: weight

        if(node%end_node == .true.) then
            weight = weight_function(x1-node%xc, node%r)
            f1 = f1+weight*data_fit_evaluate(x1, node%approx_order, node%c)
            sum_weight = sum_weight+weight
            res = 1
            return
        else
            do k = 1, node%num_child(3)
            do j = 1, node%num_child(2)
            do i = 1, node%num_child(1)
                res = subdomain_evaluate(qi, node%child(i,j,k), x1, f1, sum_weight)
            end do
            end do
            end do
        end if

        res = 2

        return
    end function subdomain_evaluate

    subroutine partition_of_unity_derivative_batch(qi, num_data, x1, l, o, d1)
        integer, intent(in) :: qi, num_data
        real, intent(in) :: x1(num_dim,num_data)
        integer, intent(in) :: l, o
        real, intent(out) :: d1(num_data)

        integer :: i

        do i = 1, num_data
            call partition_of_unity_derivative_single(qi, x1(:,i), l, o, d1(i))
        end do

        return
    end subroutine partition_of_unity_derivative_batch

    subroutine partition_of_unity_derivative_single(qi, x1, l, o, d1)
        integer, intent(in) :: qi
        real, intent(in) :: x1(num_dim)
        integer, intent(in) :: l, o
        real, intent(out) :: d1

        integer :: i, j, k, res
        real :: part1, part2
        real :: sum_weight, sum_dweight

        part1 = 0.0
        part2 = 0.0
        sum_weight = 0.0
        sum_dweight = 0.0
        do k = 1, num_tree(3)
        do j = 1, num_tree(2)
        do i = 1, num_tree(1)
            res = subdomain_derivative(qi, forest(qi)%tree(i,j,k), x1, l, o, &
                part1, part2, sum_weight, sum_dweight)
        end do
        end do
        end do
        d1 = part1/sum_weight+part2*(-1/sum_weight**2*sum_dweight)

        return
    end subroutine partition_of_unity_derivative_single

    recursive integer function subdomain_derivative(qi, node, x1, l, o, &
        part1, part2, sum_weight, sum_dweight) result(res)
        integer, intent(in) :: qi
        type(tree_node), intent(in) :: node
        real, intent(in) :: x1(num_dim)
        integer, intent(in) :: l, o
        real, intent(out) :: part1, part2
        real, intent(out) :: sum_weight
        real, intent(out) :: sum_dweight

        integer :: i, j, k
        real :: temp1, temp2
        real :: w, dw, d(num_dim)

        if(associated(node%child) == .false.) then
            d  = x1-node%xc
            w  = weight_function(d, node%r)
            dw = weight_function_derivative_1st(d, node%r, l)
            temp1 = data_fit_evaluate(x1, node%approx_order, node%c)
            temp2 = data_fit_derivative(x1, node%approx_order, node%c, l, o)
            part1 = part1+dw*temp1+w*temp2
            part2 = part2+w*temp1
            sum_weight = sum_weight+w
            sum_dweight = sum_dweight+dw
            res = 1
            return
        else
            do k = 1, node%num_child(3)
            do j = 1, node%num_child(2)
            do i = 1, node%num_child(1)
                res = subdomain_derivative(qi, node%child(i,j,k), x1, l, o, &
                    part1, part2, sum_weight, sum_dweight)
            end do
            end do
            end do
        end if

        res = 2

        return
    end function subdomain_derivative

    real function weight_function(d, c) result(w)
        real, intent(in) :: d(num_dim)
        real, intent(in) :: c(*)

        integer :: i
        real :: temp

        select case(weight_function_type)
        case("exponential")
            temp = 0.0
            do i = 1, num_dim
                temp = temp+(d(i)/c(i))**2/4
            end do
            w = exp(-exp_factor*temp)
        case default
            call message_show(ERROR_MSG, &
                "partition_of_unity",    &
                "weight_function",       &
                "invalid weight function type")
            stop
        end select

        return
    end function weight_function

    real function weight_function_derivative_1st(d, c, l) result(res)
        real, intent(in) :: d(num_dim)
        real, intent(in) :: c(*)
        integer, intent(in) :: l

        integer :: i
        real :: temp

        select case(weight_function_type)
        case("exponential")
            temp = 0.0
            do i = 1, num_dim
                temp = temp+(d(i)/c(i))**2/4
            end do
            res = exp(-exp_factor*temp)
            temp = -exp_factor/2/(c(l))**2*d(l)
            res = res*temp
        case default
            call message_show(ERROR_MSG, &
                "partition_of_unity",    &
                "weight_function",       &
                "invalid weight function type")
            stop
        end select

        return
    end function weight_function_derivative_1st

end module partition_of_unity
